import ollama
import json
import re
import regex
import requests

client = ollama.Client(host='http://localhost:11434')

def extract_json(text):
    """增强版JSON提取，支持嵌套结构"""
    try:
        # 匹配带代码块的JSON
        code_block = regex.search(r'```json\n({(?:[^{}]|(?R))*})\n```', text, regex.DOTALL)
        if code_block:
            return code_block.group(1).strip()

        # 使用递归正则匹配平衡括号
        json_match = regex.search(r'({(?:[^{}]|(?R))*})', text, regex.DOTALL)
        return json_match.group() if json_match else None
    except:
        return None

def validate_request(data):
    """验证请求合规性"""
    try:
        req = json.loads(data)
        return req.get("method") == "check" and req.get("params", {}).get("vendor") in ["all", "huawei", "cisco", "ruijie", "h3c"]
    except:
        return False

def generate_mcp_request(prompt, max_retries=3):
    """生成合规请求"""
    for _ in range(max_retries):
        try:
            response = client.generate(
                model='netbot',
                prompt=f"用户指令：{prompt}",
                format="json",
                options={'temperature': 0}  # 降低随机性
            )
            # json_str = response['response']
            json_str = extract_json(response['response'])
            print(f"提取到JSON：{json_str}")
            if json_str and validate_request(json_str):
                return json_str
        except:
            continue
    return None

# 跳过LLM引擎，直接使用规则引擎
# def generate_mcp_request(prompt):
#     """基于关键词的规则引擎"""
#     vendor_keywords = {
#         "all": ["全部", "所有", "全网", "整体"],
#         "huawei": ["华为", "Huawei", "HW","huawei","hw"],
#         "h3c": ["华三", "h3c", "H3C"],
#         "cisco": ["思科", "Cisco","cisco"],
#         "ruijie": ["锐捷", "Ruijie","ruijie","RJ",'rj']
#     }
    
#     for vendor, keywords in vendor_keywords.items():
#         if any(kw in prompt for kw in keywords):
#             return json.dumps({"method":"check","params":{"vendor": vendor}})
    
#     return json.dumps({"method":"check","params":{"vendor":"all"}})


def execute_check(request_json):
    """执行巡检请求"""
    try:
        req = json.loads(request_json)
        resp = requests.post("http://localhost:8000/mcp", json=req)
        return resp.json()
    except Exception as e:
        return {"status": "error", "message": str(e)}

if __name__ == "__main__":
    print("网络设备巡检助手（输入exit退出）")
    while True:
        cmd = input("\n请输入指令：").strip()
        if cmd.lower() in ["exit", "quit"]:
            break
        
        print("🔍 解析指令中...")
        mcp_req = generate_mcp_request(cmd)
        
        if mcp_req:
            print(f"生成请求：{mcp_req}")
            print("⚙️ 执行巡检...")
            result = execute_check(mcp_req)
            
            if result.get("status") == "success":
                with open(result["report"], 'r', encoding='utf-8') as f:
                    print("\n" + "="*40)
                    print(f.read())
                    print("="*40)
            else:
                print(f"❌ 失败：{result.get('message')}")
        else:
            print("无法解析指令，请尝试以下格式：")
            print("示例：检查所有华为设备 / 全网巡检")